from . import (
    IFileLoader,
    PDFLoader,
    TXTLoader,
    DOCXLoader,
    CSVLoader,
    EXCELLoader,
    ImageLoader,
)
from ...exception import TypeMismatchException


class DocumentLoaderFactory:
    _loaders = {
        "pdf": PDFLoader,
        "txt": TXTLoader,
        "docx": DOCXLoader,
        "csv": CSVLoader,
        "xlsx": EXCELLoader,
        "png": ImageLoader,
        "jpg": ImageLoader,
        "jpeg": ImageLoader,
    }

    @classmethod
    def get_loader(cls, file_type: str) -> IFileLoader:
        loader_cls = cls._loaders.get(file_type.lower())
        if not loader_cls:
            raise TypeMismatchException(cls._loaders.keys(), file_type)
        return loader_cls()
